import requests
import traceback
from typing import List, Dict, Any
from .db_manager import DatabaseManager
from .ai_config import AI_CONFIG, PROMPT_TEMPLATE
import os
from .ai_query_base import AIQueryBase

class AIQueryManager(AIQueryBase):
    """
    AIQueryManager 用于通过 DeepSeek API 将自然语言问题转为 SQL 并查询数据库。
    负责生成提示词、调用 AI 接口、处理 SQL 及执行数据库查询。
    """
    def __init__(self):
        """
        初始化 AIQueryManager，创建数据库管理器和请求头。
        """
        super().__init__()
        self.headers = {
            "Authorization": f"Bearer {AI_CONFIG['api_key']}",
            "Content-Type": "application/json"
        }
    

    def _get_ai_response(self, prompt: str) -> str:
        """
        调用 DeepSeek API 获取 SQL 响应。
        参数：
            prompt (str): AI prompt。
        返回：
            str: AI 返回的 SQL 语句。
        失败时返回 None。
        """
        try:
            print(f"正在调用DeepSeek API...")
            response = requests.post(
                f"{AI_CONFIG['api_base']}/chat/completions",
                headers=self.headers,
                json={
                    "model": AI_CONFIG['model'],
                    "messages": [
                        {"role": "system", "content": "请注意针对PostgreSQL数据库，请使用正确的语法。例如涉及到user等是PostgreSQL的保留关键字相关的数据库表名或字段，需要使用双引号包裹。（例如虽然是system_user，但 PostgreSQL 在某些解析情况下仍然容易对类似的名字卡壳，需要使用双引号包裹system_user）。"
                         +"请注意正确理解个字段的数据类型，避免错误数据转换。\n"},
                        {"role": "user", "content": prompt}
                    ],
                    "temperature": AI_CONFIG['temperature'],
                    "max_tokens": AI_CONFIG['max_tokens']
                }
            )
            
            if response.status_code == 200:
                result = response.json()
                content = result['choices'][0]['message']['content'].strip()
                self.log_ai_interaction(prompt, content, True)
                return content
            else:
                error_msg = f"API调用失败: {response.status_code} - {response.text}"
                print(error_msg)
                self.log_ai_interaction(prompt, response.text, False, error_msg)
                return None
                
        except Exception as e:
            error_msg = f"AI模型调用失败: {str(e)}"
            print(error_msg)
            print(traceback.format_exc())
            self.log_ai_interaction(prompt, None, False, error_msg + '\n' + traceback.format_exc())
            return None
    
    
    def get_valid_sql(self, question: str) -> (bool, str):
        """
        根据自然语言问题，多轮尝试生成并验证可执行的SQL语句。
        参数：
            question (str): 自然语言问题。
        返回：
            (bool, str): (是否成功, 最终SQL或失败原因)
        """
        max_retry = int(os.getenv('SQL_RETRY', 3))  # 可在.env中配置SQL_RETRY，默认3
        retry_count = 0
        prompt = self._generate_query_prompt(question)
        sql_query = self._get_ai_response(prompt)
        last_error = ''
        while retry_count < max_retry:
            valid, sql_checked = self.db.validate_and_fix_sql(sql_query)
            if valid:
                return True, sql_checked
            else:
                retry_count += 1
                last_error = sql_checked
                print(f"第{retry_count}次SQL验证失败，原因: {last_error}")
                prompt = self._generate_query_prompt(question) + f"\n上一次生成的SQL: {sql_query}\n验证失败原因: {last_error}\n请根据上述反馈重新生成正确的SQL语句。"
                sql_query = self._get_ai_response(prompt)
                if not sql_query:
                    break
        return False, last_error
    
    def query(self, question: str) -> Dict[str, Any]:
        """
        处理自然语言查询，自动生成 SQL 并执行。
        参数：
            question (str): 自然语言问题。
        返回：
            dict: 查询结果，包含 SQL、执行结果或错误信息。
        """
        success, sql_or_err = self.get_valid_sql(question)
        if success:
            print(f"处理后的SQL: {sql_or_err}")
            try:
                results = self.db.execute_query(sql_or_err)
                return {
                    'success': True,
                    'sql': sql_or_err,
                    'results': results
                }
            except Exception as e:
                return {
                    'success': False,
                    'error': f'SQL执行失败: {str(e)}',
                    'traceback': traceback.format_exc()
                }
        else:
            return {
                'success': False,
                'error': f'SQL验证失败，重试后仍未通过: {sql_or_err}',
                'traceback': traceback.format_exc()
            }